# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/5 14:29 
@Author  : ZhangShenao 
@File    : 2.使用FileChatMessageHistory组件保存聊天历史.py
@Desc    : 使用FileChatMessageHistory组件,将聊天历史保存到文件中
"""
import os

import dotenv
from langchain_community.chat_message_histories import FileChatMessageHistory
from langchain_core.messages import AIMessage, HumanMessage
from zhipuai import ZhipuAI

# 系统指令
SYSTEM_PROMPT_TEMPLATE = '''你是一个具有记忆功能的AI助手，请根据结合用户的当前提问，和之前的历史聊天记录，生成准确的回答。你们的对话全部采用中文。
下面是你们的历史聊天记录：
<历史聊天记录开始>

{chat_history}

<历史聊天记录结束>
'''

# 调用的模型名称
MODEL_NAME = 'glm-4-air'

# 加载环境变量
dotenv.load_dotenv()

# 创建智普AI客户端
client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))

# 创建FileChatMessageHistory组件,将聊天历史保存到文件中
chat_message_history = FileChatMessageHistory(file_path='./chat_history.json')

# 接收用户输入,调用LLM
while True:
    query = input('Human：')
    if 'bye' == query:
        print('bye bye~')
        break

    # 填充历史聊天记录
    # 因为BaseChatMessageHistory重写了__str__函数,所以可以直接将其作为str参数传入
    system_prompt = SYSTEM_PROMPT_TEMPLATE.format(chat_history=chat_message_history)

    # 流式调用LLM
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {'role': 'system', 'content': system_prompt},
            {"role": "user", "content": query},
        ],
        stream=True,
    )

    # 获取LLM输出结果
    output = ''
    print('AI：', end='', flush=True)
    for chunk in response:
        content = chunk.choices[0].delta.content
        print(content, end='', flush=True)
        output += content
    print()

    # 将用户和AI的对话记录保存到历史聊天记录中
    chat_message_history.add_messages([HumanMessage(query), AIMessage(output)])
